import torch
import torchvision
from thop import profile
from models.bifovnet import BiFovNet


x1 = torch.randn(1, 3, 224, 224).cuda()
x2 = torch.randn(1, 3, 224, 224).cuda()#.cuda()

model = BiFovNet(dim=48, depth=[3, 3, 12, 5], in_chans=3, kernel_size=7, patch_size=4,
                   num_classes=16, H=224, W=224, p_h=[8, 4, 2, 1], p_w=[8, 4, 2, 1])

model.cuda()
model.eval()
flops, params = profile(model, inputs=(x1, x2))
# flops, params = profile(model, inputs=(x1))
print("Flops:", "%.2fM" % (flops / 1e6), "Params:", "%.2fM" % (params / 1e6))
